import sys
import os
import cv2
import threading
from PySide6.QtWidgets import QApplication, QWidget
from PySide6.QtGui import QImage, QPixmap
from PySide6.QtCore import QTimer, Qt
from ultralytics import YOLO
from ui_form import Ui_Widget

class ImageDetectionWindow(QWidget):
    def __init__(self):
        super().__init__()
        # 初始化 UI
        self.ui = Ui_Widget()
        self.ui.setupUi(self)

        # 加载 YOLOv8 模型
        self.model = YOLO(r"D:\jupyterCode\ship\runs\detect\train\weights\best.pt")


        # 配置
        self.image_folder = "D:/qt/qtproject/untitled/picture"  # 输入图片文件夹
        self.output_folder = "D:/qt/qtproject/untitled/output"  # 输出结果文件夹
        self.image_list = []
        self.current_index = 0

        # 创建输出文件夹
        if not os.path.exists(self.output_folder):
            os.makedirs(self.output_folder)

        # 设置定时器，用于逐张处理图片
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.process_next_image)

        # 连接按钮信号
        self.ui.startButton.clicked.connect(self.start_detection)

    def start_detection(self):
        # 加载图片列表
        self.image_list = [f for f in os.listdir(self.image_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]
        if not self.image_list:
            print("文件夹中没有图片！")
            return

        self.current_index = 0
        self.ui.progressBar.setMaximum(len(self.image_list))  # 设置进度条最大值
        self.ui.progressBar.setValue(0)  # 初始化进度条
        self.timer.start(100)  # 每 100ms 处理一张图（可调整速度）

    def process_next_image(self):
        if self.current_index < len(self.image_list):
            # 读取图片
            image_path = os.path.join(self.image_folder, self.image_list[self.current_index])
            frame = cv2.imread(image_path)

            if frame is not None:
                # 使用 YOLOv8 进行目标检测
                results = self.model(frame)
                annotated_frame = results[0].plot()

                # 保存检测结果
                output_path = os.path.join(self.output_folder, f"detected_{self.image_list[self.current_index]}")
                cv2.imwrite(output_path, annotated_frame)

                # 显示当前图片
                frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
                h, w, ch = frame_rgb.shape
                bytes_per_line = ch * w
                qt_image = QImage(frame_rgb.data, w, h, bytes_per_line, QImage.Format_RGB888)
                pixmap = QPixmap.fromImage(qt_image)
                self.ui.videoLabel.setPixmap(pixmap.scaled(self.ui.videoLabel.size(), aspectMode=Qt.KeepAspectRatio))

                # 更新进度条
                self.ui.progressBar.setValue(self.current_index + 1)
                print(f"已处理: {self.current_index + 1}/{len(self.image_list)} - {self.image_list[self.current_index]}")

            self.current_index += 1
        else:
            # 处理完成
            self.timer.stop()
            print("所有图片处理完成！")
            self.ui.videoLabel.clear()

    def closeEvent(self, event):
        self.timer.stop()
        event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ImageDetectionWindow()
    window.show()
    sys.exit(app.exec())
